//=========================================================================
// Copyright (C) 2024 The C++ Component Model(COMO) Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//=========================================================================

/*  =========================================================================
    zmtp_channel - channel class

    Copyright (c) contributors as noted in the AUTHORS file.
    This file is part of libzmtp, the C ZMTP stack.

    This Source Code Form is subject to the terms of the Mozilla Public
    License, v. 2.0. If a copy of the MPL was not distributed with this
    file, You can obtain one at http://mozilla.org/MPL/2.0/.
    =========================================================================
*/

#include "zmtp_classes.h"

// ZMTP greeting (64 bytes)
#pragma pack(1)
struct zmtp_greeting {
    byte signature[10];
    byte version[2];
    byte mechanism[20];
    byte as_server[1];
    byte filler[31];
};
#pragma pack()

static zmtp_endpoint_t *s_endpoint_from_str(const char *endpoint_str);
static int s_negotiate(zmtp_channel_t *self);
static int s_tcp_send(int fd, const void *data, size_t len);
static int s_tcp_recv(int fd, void *buffer, size_t len);

CMemPool *g_zmtp_channel_buffer = NULL;

/**
 * Constructor
 */
zmtp_channel_t *zmtp_channel_new()
{
    zmtp_channel_t *self = NULL;
    if (NULL != g_zmtp_channel_buffer) {
        self = (zmtp_channel_t *)CMemPool_Alloc(g_zmtp_channel_buffer,
                                     sizeof(zmtp_channel_t), MUST_USE_MEM_POOL);
    }
    if (NULL == self) {
        self = (zmtp_channel_t *)zmalloc(sizeof(zmtp_channel_t));
    }

    if (NULL == self) {
        return NULL;
    }
    self->fd = -1;
    return self;
}

/**
 * Destructor; closes fd if connected
 */
void zmtp_channel_destroy(zmtp_channel_t **self_p)
{
    if (NULL == self_p) {
        return;
    }

    if (NULL != *self_p) {
        zmtp_channel_t *self = *self_p;
        if (self->fd != -1) {
            if (self->zmtpDealerId < 0) {
                close(self->fd);
            }
        }

        if ((NULL != g_zmtp_channel_buffer) &&
                           (CMemPool_CheckExist(g_zmtp_channel_buffer, self))) {
            CMemPool_Free(g_zmtp_channel_buffer, self);
        }
        else {
            zfree(self);
        }
        *self_p = NULL;
    }
}

/**
 * Connect channel to local endpoint
 */
int zmtp_channel_ipc_connect(zmtp_channel_t *self, const char *path)
{
    if ((NULL == self) || (NULL == path)) {
        return -1;
    }

    zmtp_endpoint_t *endpoint = (zmtp_endpoint_t *)zmtp_ipc_endpoint_new(path);
    if (NULL == endpoint) {
        return -1;
    }

    self->fd = zmtp_endpoint_connect(endpoint);
    zmtp_endpoint_destroy(&endpoint);
    if (-1 == self->fd) {
        return -1;
    }

    if (s_negotiate(self) < 0) {
        if (self->zmtpDealerId < 0) {
            close(self->fd);
        }
        self->fd = -1;
        return -1;
    }

    return 0;
}

/**
 * Connect channel to TCP endpoint
 */
int zmtp_channel_tcp_connect(zmtp_channel_t *self,
                                           const char *addr, unsigned short port)
{
    if ((NULL == self) || (self->fd != -1)) {
        return -1;
    }

    zmtp_endpoint_t *endpoint = (zmtp_endpoint_t *)zmtp_tcp_endpoint_new(addr, port);
    if (NULL == endpoint) {
        return -1;
    }

    self->fd = zmtp_endpoint_connect(endpoint);
    zmtp_endpoint_destroy(&endpoint);
    if (-1 == self->fd) {
        return -1;
    }

    if (s_negotiate(self) < 0) {
        if (self->zmtpDealerId < 0) {
            close(self->fd);
        }
        self->fd = -1;
        return -1;
    }

    return 0;
}

/**
 * zmtp_channel_connect
 */
int zmtp_channel_connect(zmtp_channel_t *self, const char *endpoint_str)
{
    if ((NULL == self) || (self->fd != -1)) {
        return -1;
    }

    zmtp_endpoint_t *endpoint = s_endpoint_from_str(endpoint_str);
    if (NULL == endpoint) {
        return -2;
    }

    self->fd = zmtp_endpoint_connect(endpoint);
    zmtp_endpoint_destroy(&endpoint);
    if (-1 == self->fd) {
        return -3;
    }

    if (s_negotiate(self) < 0) {
        if (self->zmtpDealerId < 0) {
            close(self->fd);
        }
        self->fd = -1;
        return -4;
    }

    return 0;
}

/**
 * zmtp_channel_listen
 *
 * Use fd to associate existing socket, fd must be in the wait listen state.
 */
int zmtp_channel_listen(zmtp_channel_t *self, const char *endpoint_str, void *channel)
{
    if ((NULL == self) || (self->fd != -1)) {
        // (self->fd != -1) means:
        // channel has been associated with one socket and cannot be associated
        // with any other socket
        return -1;
    }

    if (channel != NULL) {
        self->fd = ((zmtp_channel_t *)channel)->fd;
        self->zmtpDealerId = ((zmtp_channel_t *)channel)->zmtpDealerId;

        if (self->zmtpDealerId < 0) {
            return -2;
        }
    }
    else {
        zmtp_endpoint_t *endpoint = s_endpoint_from_str(endpoint_str);
        if (NULL == endpoint) {
            return -3;
        }

        self->fd = zmtp_endpoint_listen(endpoint);
        self->zmtpDealerId = -1;

        zmtp_endpoint_destroy(&endpoint);
        if (-1 == self->fd) {
            return -4;
        }
    }

    if (s_negotiate(self) < 0) {
        if (self->zmtpDealerId < 0) {
            close(self->fd);
        }
        self->fd = -1;
        return -5;
    }
    return 0;
}

/**
 * s_endpoint_from_str
 */
static zmtp_endpoint_t *s_endpoint_from_str(const char *endpoint_str)
{
    if (strncmp(endpoint_str, "ipc://", 6) == 0) {
        return (zmtp_endpoint_t *)zmtp_ipc_endpoint_new(endpoint_str + 6);
    }
    else if (strncmp(endpoint_str, "tcp://", 6) == 0) {
        char *colon = strrchr(endpoint_str + 6, ':');
        if (NULL == colon) {
            return NULL;
        }
        else {
            const size_t addr_len = colon - endpoint_str - 6;
            char addr[addr_len + 1];

            memcpy(addr, endpoint_str + 6, addr_len);
            addr[addr_len] = '\0';

            const unsigned short port = atoi(colon + 1);
            return (zmtp_endpoint_t *)zmtp_tcp_endpoint_new(addr, port);
        }
    }
    else
        return NULL;
}

/**
 * Negotiate a ZMTP channel
 * This currently does only ZMTP v3, and will reject older protocols.
 * TODO: test sending random/wrong data to this handler.
 */
static int s_negotiate(zmtp_channel_t *self)
{
    if ((NULL == self) || (self->fd < 0)) {
        return -1;
    }

    int s = self->fd;

    // This is our greeting (64 octets)
    static const struct zmtp_greeting outgoing = {
        .signature = { 0xffu, 0, 0, 0, 0, 0, 0, 0, 1u, 0x7fu },
        .version   = { 3u, 0 },
        .mechanism = { 'N', 'U', 'L', 'L', '\0' }
    };

    do {
        // Send protocol signature
        if (self->zmtpDealerId < 0) {
            /**
             * Sharing someone else's channel, someone else has already sent the
             * same message to the client.
             *
             * Send protocol signature
             */
            if (s_tcp_send(s, outgoing.signature, sizeof(outgoing.signature)) == -1) {
                break;
            }
        }

        // Read the first byte.
        struct zmtp_greeting incoming = {0};
        incoming.version[0] = 0;

        if (s_tcp_recv(s, incoming.signature, 1) == -1) {
            break;
        }
        if (0xffu != incoming.signature[0]) {
            break;
        }

        // Read the rest of signature
        if (s_tcp_recv(s, incoming.signature + 1, 9) == -1) {
            break;
        }
        if (1u != (incoming.signature[9] & 1u)) {
            break;
        }

        // Exchange major version numbers
        if (s_tcp_send(s, outgoing.version, 1) == -1) {
            break;
        }

        if (s_tcp_recv(s, incoming.version, 1) == -1) {
            break;
        }

        if (3 != incoming.version[0]) {
            break;
        }

        // Send the rest of greeting to the peer.
        if (s_tcp_send(s, outgoing.version + 1, 1) == -1) {
            break;
        }
        if (s_tcp_send(s, outgoing.mechanism, sizeof(outgoing.mechanism)) == -1) {
            break;
        }
        if (s_tcp_send(s, outgoing.as_server, sizeof(outgoing.as_server)) == -1) {
            break;
        }
        if (s_tcp_send(s, outgoing.filler, sizeof(outgoing.filler)) == -1) {
            break;
        }

        // Receive the rest of greeting from the peer.
        if (s_tcp_recv(s, incoming.version + 1, 1) == -1) {
            break;
        }
        if (s_tcp_recv(s, incoming.mechanism, sizeof(incoming.mechanism)) == -1) {
            break;
        }
        if (s_tcp_recv(s, incoming.as_server, sizeof(incoming.as_server)) == -1) {
            break;
        }
        if (s_tcp_recv(s, incoming.filler, sizeof(incoming.filler)) == -1) {
            break;
        }

        // Send READY command
        zmtp_msg_t *ready = zmtp_msg_from_const_data(0x04, "\5READY", 6);
        if (NULL == ready) {
            break;
        }

        zmtp_channel_send(self, ready);
        zmtp_msg_destroy(&ready);

        // Receive READY command
        ready = zmtp_channel_recv(self);
        if (NULL == ready) {
            break;
        }

        if ((zmtp_msg_flags(ready) & ZMTP_MSG_COMMAND) != ZMTP_MSG_COMMAND) {
            break;
        }
        zmtp_msg_destroy(&ready);

        return 0;
    } while (false);

    return -2;
}

/**
 * Send a ZMTP message to the channel
 */
int zmtp_channel_send(zmtp_channel_t *self, zmtp_msg_t *msg)
{
    if ((NULL == self) || (NULL == msg)) {
        return -1;
    }

    byte frame_flags = 0;
    if ((zmtp_msg_flags(msg) & ZMTP_MSG_MORE) == ZMTP_MSG_MORE) {
        frame_flags |= ZMTP_MORE_FLAG;
    }
    if ((zmtp_msg_flags(msg) & ZMTP_MSG_COMMAND) == ZMTP_MSG_COMMAND) {
        frame_flags |= ZMTP_COMMAND_FLAG;
    }
    if (zmtp_msg_size(msg) > 255) {
        frame_flags |= ZMTP_LARGE_FLAG;
    }
    if (s_tcp_send(self->fd, &frame_flags, sizeof(frame_flags)) == -1) {
        return -1;
    }

    if (zmtp_msg_size (msg) <= 255) {
        const byte msg_size = zmtp_msg_size (msg);
        if (s_tcp_send (self->fd, &msg_size, sizeof(msg_size)) == -1) {
            return -1;
        }
    }
    else {
        byte buffer[8];
        const uint64_t msg_size = (uint64_t)zmtp_msg_size(msg);
        buffer[0] = msg_size >> 56;
        buffer[1] = msg_size >> 48;
        buffer[2] = msg_size >> 40;
        buffer[3] = msg_size >> 32;
        buffer[4] = msg_size >> 24;
        buffer[5] = msg_size >> 16;
        buffer[6] = msg_size >> 8;
        buffer[7] = msg_size;
        if (s_tcp_send(self->fd, buffer, sizeof(buffer)) == -1) {
            return -1;
        }
    }
    if (s_tcp_send(self->fd, zmtp_msg_data(msg), zmtp_msg_size(msg)) == -1) {
        return -1;
    }
    return 0;
}

/**
 * Receive a ZMTP message off the channel
 */
zmtp_msg_t *zmtp_channel_recv(zmtp_channel_t *self)
{
    if (NULL == self) {
        return NULL;
    }

    byte frame_flags;
    size_t size;

    if (s_tcp_recv(self->fd, &frame_flags, 1) == -1) {
        return NULL;
    }

    // Check large flag
    if ((frame_flags & ZMTP_LARGE_FLAG) == 0) {
        byte buffer[1];
        if (s_tcp_recv(self->fd, buffer, 1) == -1) {
            return NULL;
        }
        size = (size_t)buffer[0];
    }
    else {
        byte buffer[8];
        if (s_tcp_recv(self->fd, buffer, sizeof(buffer)) == -1) {
            return NULL;
        }
        size = (uint64_t)buffer[0] << 56 |
               (uint64_t)buffer[1] << 48 |
               (uint64_t)buffer[2] << 40 |
               (uint64_t)buffer[3] << 32 |
               (uint64_t)buffer[4] << 24 |
               (uint64_t)buffer[5] << 16 |
               (uint64_t)buffer[6] << 8  |
               (uint64_t)buffer[7];
    }

    byte *data = zmalloc(size);
    if (NULL == data) {
        return NULL;
    }

    if (s_tcp_recv(self->fd, data, size) == -1) {
        zfree(data);
        return NULL;
    }

    byte msg_flags = 0;
    if ((frame_flags & ZMTP_MORE_FLAG) == ZMTP_MORE_FLAG) {
        msg_flags |= ZMTP_MSG_MORE;
    }
    if ((frame_flags & ZMTP_COMMAND_FLAG) == ZMTP_COMMAND_FLAG) {
        msg_flags |= ZMTP_MSG_COMMAND;
    }

    return zmtp_msg_from_data(msg_flags, &data, size);
}

/**
 * Lower-level TCP and ZMTP message I/O functions
 */
static int s_tcp_send(int fd, const void *data, size_t len)
{
    size_t bytes_sent = 0;

    while (bytes_sent < len) {
        ssize_t rc;

        rc = send(fd, (char *)data + bytes_sent, len - bytes_sent, 0);
        if ((rc == -1) && ((errno == EINTR) || (errno == EWOULDBLOCK) ||
                                                           (errno == EAGAIN))) {
            continue;
        }
        if (-1 == rc) {
            return -1;
        }
        bytes_sent += rc;
    }

    return 0;
}

#ifdef __SELF_TEST__
//for test
void arrayToASCII(const char* arr, int len)
{
    for (int i = 0; i < len ; i++) {
        printf("%d ", arr[i]);
    }
}
#endif

/**
 * s_tcp_recv
 */
static int s_tcp_recv(int fd, void *buffer, size_t len)
{
    size_t bytes_read = 0;
    while (bytes_read < len) {
        ssize_t n;

        n = recv(fd, (char *)buffer + bytes_read, len - bytes_read, 0);
        if ((-1 == n) && ((errno == EINTR) || (errno == EWOULDBLOCK) ||
                                                           (errno == EAGAIN))) {
            continue;
        }
        if ((-1 == n) || (0 == n)) {
            return -1;
        }
        bytes_read += n;
    }

#ifdef __SELF_TEST__
    printf("Read message, ascii: ["); // for test
    arrayToASCII(buffer, len);
    printf("]\n"); // for test
#endif

    return 0;
}

#ifdef __SELF_TEST__

#include <poll.h>

// Simple TCP echo server. It listens on a TCP port and after
// accepting a new connection, echoes all received data.
// This is to test the encodining/decoding compatibility.

struct echo_serv_t {
    unsigned short port;
};

static void *s_echo_serv(void *arg)
{
    struct echo_serv_t *params = (struct echo_serv_t *)arg;

    // Create socket
    const int s = socket(AF_INET, SOCK_STREAM, 0);
    if (-1 == s) {
        return NULL;
    }

    // Allow port reuse
    const int on = 1;
    int rc = setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
    if (0 != rc) {
        return NULL;
    }

    // Fill address
    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(params->port);
    server_addr.sin_addr.s_addr = htonl(INADDR_ANY);

    // Bind socket
    rc = bind(s, (struct sockaddr *)&server_addr, sizeof(server_addr));
    if (0 != rc) {
        return NULL;
    }

    // Listen for connections
    rc = listen(s, 1);
    if (-1 == rc) {
        return NULL;
    }

    // Accept connection
    int fd = accept(s, NULL, NULL);
    if (-1 == fd) {
        return NULL;
    }

    // Set non-blocking mode
    const int flags = fcntl(fd, F_GETFL, 0);
    if (-1 == flags) {
        return NULL;
    }

    rc = fcntl(fd, F_SETFL, flags | O_NONBLOCK);
    if (0 != rc) {
        return NULL;
    }

    unsigned char buf[80];

    // Echo all received data
    while (1) {
        struct pollfd pollfd;
        pollfd.fd = fd;
        pollfd.events = POLLIN;
        rc = poll(&pollfd, 1, -1);
        if (1 != rc) {
            return NULL;
        }

        rc = read(fd, buf, sizeof(buf));
        if (0 == rc) {
            break;
        }
        assert (rc > 0 || errno == EINTR);
        if (rc > 0) {
            rc = s_tcp_send(fd, buf, rc);
            if (0 != rc) {
                return NULL;
            }
        }
    }
    close(fd);
    close(s);

    return NULL;
}

struct script_line {
    char cmd;           // 'i' for input, 'o' for output, 'x' terminator
    size_t data_len;    // length of data
    const char *data;   // data to send or expect
};

struct test_server_t {
    unsigned short port;
    const struct script_line *script;
};

static void *s_test_server(void *arg)
{
    struct test_server_t *params = (struct test_server_t *)arg;

    // Create socket
    const int s = socket(AF_INET, SOCK_STREAM, 0);
    if (-1 == s) {
        return NULL;
    }

    // Allow port reuse
    const int on = 1;
    int rc = setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
    if (0 != rc) {
        return NULL;
    }

    // Fill address
    struct sockaddr_in server_addr = {
        .sin_family = AF_INET,
        .sin_port = htons(params->port),
        .sin_addr.s_addr = htonl(INADDR_ANY),
    };
    // Bind socket
    rc = bind(s, (struct sockaddr *)&server_addr, sizeof(server_addr));
    if (0 != rc) {
        return NULL;
    }

    // Listen for connections
    rc = listen(s, 1);
    if (-1 == rc) {
        return NULL;
    }

    // Accept connection
    const int fd = accept(s, NULL, NULL);
    if (-1 == fd) {
        return NULL;
    }

    // Run I/O script
    for (int i = 0;  params->script[i].cmd != 'x';  i++) {
        const char cmd = params->script[i].cmd;
        const size_t data_len = params->script[i].data_len;
        const char *data = params->script[i].data;
        assert (cmd == 'i' || cmd == 'o');

        if ('i' == cmd) {
            char buf[data_len];
            const int rc = s_tcp_recv(fd, buf, data_len);
            assert (0 == rc);
            assert (memcmp(buf, data, data_len) == 0);
        }
        else {
            const int rc = s_tcp_send(fd, data, data_len);
            assert (0 == rc);
        }
    }

    close(fd);
    close(s);
    return NULL;
}

// --------------------------------------------------------------------------
// Selftest

void zmtp_channel_test(bool verbose)
{
    printf (" * zmtp_channel: ");
    // @selftest
    pthread_t thread;

    struct echo_serv_t echo_serv_params = { .port = 22001 };

    pthread_create(&thread, NULL, s_echo_serv, &echo_serv_params);
    sleep(1);

    zmtp_channel_t *channel = zmtp_channel_new();
    assert (channel);

    int rc = zmtp_channel_tcp_connect(channel, "127.0.0.1", 22001);
    assert (0 == rc);

    char *test_strings[] = {
        "1",
        "22",
        "333",
        "4444",
        "55555"
    };

    for (int i = 0;  i < 5;  i++) {
        zmtp_msg_t *msg = zmtp_msg_from_const_data(0, test_strings[i],
                                                   strlen(test_strings[i]));
        assert (msg);
        rc = zmtp_channel_send(channel, msg);
        assert (0 == rc);
        zmtp_msg_t *msg2 = zmtp_channel_recv(channel);
        assert (msg2 != NULL);
        assert (zmtp_msg_size(msg) == zmtp_msg_size(msg2));
        assert (memcmp(zmtp_msg_data(msg),
                        zmtp_msg_data(msg2), zmtp_msg_size(msg)) == 0);

        zmtp_msg_destroy(&msg);
        zmtp_msg_destroy(&msg2);
    }

    zmtp_channel_destroy(&channel);
    pthread_join(thread, NULL);

    // Test flow, initial handshake, receive "ping 1" and "ping 2" messages,
    // then send "pong 1" and "ping 2"
    struct script_line script[] = {
        { 'o', 10, "\xFF\0\0\0\0\0\0\0\1\x7F" },
        { 'o', 2, "\3\0" },
        { 'o', 20, "NULL\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0" },
        { 'o', 32, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0" },
        { 'i', 10, "\xFF\0\0\0\0\0\0\0\1\x7F" },
        { 'i', 2, "\3\0" },
        { 'i', 20, "NULL\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0" },
        { 'i', 32, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0" },
        { 'o', 8, "\4\6\5READY" },     // send READY command
        { 'i', 8, "\4\6\5READY" },     // expect READY command
        { 'i', 8, "\1\6ping 1" },      // expect ping 1, more set
        { 'i', 8, "\0\6ping 2" },      // expect ping 2, more flag not set
        { 'o', 8, "\1\6pong 1" },      // send pong 1, more set
        { 'o', 8, "\0\6pong 2" },      // send pong 2, more flag not set
        { 'x' },
    };

    struct test_server_t params = {
        .port = 22000,
        .script = script,
    };
    pthread_create(&thread, NULL, s_test_server, &params);
    sleep (1);

    channel = zmtp_channel_new();
    assert (channel);
    rc = zmtp_channel_tcp_connect(channel, "127.0.0.1", 22000);
    assert (0 == rc);

    // Send "ping 1"
    zmtp_msg_t *ping_1 = zmtp_msg_from_const_data(ZMTP_MSG_MORE, "ping 1", 6);
    rc = zmtp_channel_send(channel, ping_1);
    assert (0 == rc);
    zmtp_msg_destroy(&ping_1);

    // Send "ping 2"
    zmtp_msg_t *ping_2 = zmtp_msg_from_const_data(0, "ping 2", 6);
    rc = zmtp_channel_send(channel, ping_2);
    assert (0 == rc);
    zmtp_msg_destroy(&ping_2);

    // Receive "pong 1"
    zmtp_msg_t *pong_1 = zmtp_channel_recv(channel);
    assert (pong_1 != NULL);
    assert (zmtp_msg_size(pong_1) == 6);
    assert (memcmp (zmtp_msg_data(pong_1), "pong 1", 6) == 0);
    assert ((zmtp_msg_flags(pong_1) & ZMTP_MSG_MORE) == ZMTP_MSG_MORE);
    zmtp_msg_destroy(&pong_1);

    // Receive "pong 2"
    zmtp_msg_t *pong_2 = zmtp_channel_recv(channel);
    assert (pong_2 != NULL);
    assert (zmtp_msg_size(pong_2) == 6);
    assert (memcmp(zmtp_msg_data(pong_2), "pong 2", 6) == 0);
    assert ((zmtp_msg_flags(pong_2) & ZMTP_MSG_MORE) == 0);

    zmtp_msg_destroy(&pong_2);

    zmtp_channel_destroy(&channel);
    pthread_join(thread, NULL);

    // @end
    printf("OK\n");
}
#endif
